import os
from pycocotools.coco import COCO
import argparse
from PIL import Image, ImageDraw
from tqdm import tqdm


def load_det_results(coco: COCO, det_file) -> COCO:
    coco_det = coco.loadRes(det_file)
    coco_det.createIndex()
    return coco_det


def load_COCO_info(ann_file):
    coco = COCO(ann_file)
    cat_ids = coco.getCatIds()
    cats = coco.loadCats(cat_ids)
    cat_names = {cat['id']: cat['name'] for cat in cats}
    return coco, cat_names


def get_args():
    parser = argparse.ArgumentParser(description="Plot json results given by object detection model.")
    parser.add_argument("--ann_file", type=str, help="Path to dataset ground truth annotation file.")
    parser.add_argument("--det_file", type=str, help="Path to json results given by object detection model.")
    parser.add_argument("--src_img_dir", type=str, help="Directory path to images which will be plotted with bboxes.")
    parser.add_argument("--dst_img_dir", type=str, help="Directory path where to saves the images with bboxes.")
    parser.add_argument("--threshold", type=float, default=0.5, help="Score threshold to filter the bboxes.")
    args = parser.parse_args()
    return args

def plot_bboxes():
    args = get_args()
    coco, coco_names = load_COCO_info(args.ann_file)
    coco_det = load_det_results(coco, args.det_file)
    det_anns = coco_det.imgToAnns
    det_img_keys = list(det_anns.keys())
    coco_imgs_info = coco.loadImgs(det_img_keys)
    for img_info in tqdm(coco_imgs_info):
        file_name = img_info['file_name']
        file_id = img_info['id']
        src_path = os.path.join(args.src_img_dir, file_name)
        dst_path = os.path.join(args.dst_img_dir, file_name)
        img = Image.open(src_path).convert("RGB")
        img_anns = det_anns[file_id]
        img_draw = ImageDraw.Draw(img)
        for ann in img_anns:
            score = ann['score']
            if score < args.threshold:
                continue
            x, y, w, h = ann['bbox']
            cat_id = ann['category_id']
            img_draw.rectangle([(x, y), (x + w, y + h)], outline="green")    # plot bbox
            img_draw.text((x, max(0, y - 10)), coco_names[cat_id])
        img.save(dst_path, "JPEG")


if __name__ == "__main__":
    plot_bboxes()
